import torch
import torch.nn as nn
from torchvision import models, datasets
from torchvision.models import MobileNet_V3_Large_Weights
import os

# 定义数据集路径
data_dir = 'Nonsegmented'  # 替换为你的数据集路径

# 使用 ImageFolder 加载数据集以确定类别名称
dataset = datasets.ImageFolder(data_dir)
class_names = dataset.classes  # 自动获取类别名称
print(f"Detected classes: {class_names}")

# 构建与训练时相同的模型结构
model = models.mobilenet_v3_large(weights=MobileNet_V3_Large_Weights.IMAGENET1K_V1)
num_ftrs = model.classifier[3].in_features
model.classifier[3] = nn.Linear(num_ftrs, len(class_names))  # 修改最后一层以匹配类别数

# 加载训练得到的权重
model.load_state_dict(torch.load('./mobilenet/best_model.pth'))

# 设置为评估模式
model.eval()

# 随机生成一个输入张量（输入形状需要与你训练时一致，例如224x224的3通道图像）
dummy_input = torch.randn(1, 3, 224, 224)

# 转换为 ONNX 格式
torch.onnx.export(
    model,                    # 模型
    dummy_input,               # 输入张量
    "mobilenet_v3_large.onnx", # 输出 ONNX 文件名
    export_params=True,        # 是否导出模型参数
    opset_version=11,          # ONNX opset 版本
    do_constant_folding=True,  # 是否进行常量折叠优化
    input_names=['input'],     # 输入节点名称
    output_names=['output'],   # 输出节点名称
    dynamic_axes={             # 动态维度设置，若需要动态输入尺寸可以设置
        'input': {0: 'batch_size'}, 
        'output': {0: 'batch_size'}
    }
)

print(f"ONNX 模型已成功导出到 mobilenet_v3_large.onnx")
